"""
使用**市赚率**来挑选股票。

市赚率（PR） = 市盈率（PE） / ROE

若：市赚率 < 1，说明被低估

若是市赚率 <= 0.5，则说明被严重低估！
"""
from utils.methods import *
import os
import baostock as bs

MIN_YEAR= 2000
START_YEAR= 2020
END_YEAR = 2025
END_DATE = '2025-10-20'
HS300_FILE_PATH = os.path.expanduser(f'~/Desktop/hs300_data_{END_YEAR}.csv')
SZ50_FILE_PATH = os.path.expanduser(f'~/Desktop/sz50_data_{END_YEAR}.csv')
ALL_STOCK_FILE_PATH = os.path.expanduser(f'~/Desktop/all_stock_data_{END_YEAR}.csv')
STABLE_CODE_FILE_PATH = os.path.expanduser(f'~/Desktop/stable_growth_{END_YEAR}.csv')
NUM_THREADS = 10


def is_stock(row_data):
    """
    判断是否是股票
    :param row_data: 股票数据行
    :return: 是否是股票
    """
    if '指数' in row_data["code_name"]:
        return False
    if '上证' in row_data["code_name"]:
        return False
    return True



'''
输出策略1的股票
'''
if __name__ == '__main__':
    # # 获取沪深300数据
    # all_data = None
    # if os.path.isfile(HS300_FILE_PATH):
    #     print(f"{HS300_FILE_PATH} 存在")
    #     all_data = pd.read_csv(HS300_FILE_PATH)
    # else:
    #     bs.login()
    #     all_data = get_hs300_data(END_DATE)
    #     all_data.to_csv(HS300_FILE_PATH, index=False)
    #     print(all_data)
    #     bs.logout()

    # 获取上证50数据
    # all_data = None
    # if os.path.isfile(SZ50_FILE_PATH):
    #     print(f"{SZ50_FILE_PATH} 存在")
    #     all_data = pd.read_csv(SZ50_FILE_PATH)
    # else:
    #     bs.login()
    #     all_data = get_sh_50_data(END_DATE)
    #     all_data.to_csv(SZ50_FILE_PATH, index=False)
    #     print(all_data)
    #     bs.logout()

    # 获取所有数据
    all_data = None
    if os.path.isfile(ALL_STOCK_FILE_PATH):
        print(f"{ALL_STOCK_FILE_PATH} 存在")
        all_data = pd.read_csv(ALL_STOCK_FILE_PATH)
    else:
        bs.login()
        all_data = get_all_stock_data(END_DATE)
        all_data.to_csv(ALL_STOCK_FILE_PATH, index=False)
        print(all_data)
        bs.logout()

    # 过滤业绩稳定的公司
    stable_code_list = []
    if os.path.isfile(STABLE_CODE_FILE_PATH):
        print(f"{STABLE_CODE_FILE_PATH} 存在")
        stable_code_list = pd.read_csv(STABLE_CODE_FILE_PATH)
    else:
        bs.login()
        i = 1
        for index, row in all_data.iterrows():
            if not is_stock(row):
                continue
            # 判断股票的盈利能力是否稳定
            base_str = f'{i}. {row["code"]} {row["code_name"]} is '
            if is_profit_stable(row['code'], MIN_YEAR, END_YEAR):
                print(base_str + 'stable')
                stable_code_list.append(row)
            else:
                print(base_str + 'NOT stable')
            i += 1
        bs.logout()
        stable_code_list = pd.DataFrame(stable_code_list)
        stable_code_list.to_csv(STABLE_CODE_FILE_PATH, index=False)

    # 计算市赚率，按市赚率升序排序
    result = pd.DataFrame(columns=['code', 'code_name', 'pr', 'pe', 'avg_roe'])
    bs.login()
    for index, stable_code in stable_code_list.iterrows():
        code = stable_code['code']
        try:
            pe = get_pe(code, END_DATE)
            if pe < 0:
                print(f'{index + 1}/{len(stable_code_list)}. {code} {stable_code['code_name']} skipped: pe < 0.')
                continue
            avg_roe = get_roe(code, START_YEAR, END_YEAR)
            result.loc[len(result)] = [
                code,
                stable_code['code_name'],
                round(pe / (avg_roe * 100), 4),
                round(pe, 4),
                round(avg_roe, 4)
            ]
            print(f'{index + 1}/{len(stable_code_list)}. {code} {stable_code['code_name']} done.')
        except Exception as e:
            print(e)
            print(f'{code} 计算失败')
    bs.logout()
    result = result.sort_values(by=['pr'], ascending=True)
    file_path = f"~/Desktop/low_pr_{END_DATE}.xlsx"
    save_as_excel(result, file_path)
    print(f'处理完成：{file_path}')

